import torch
from torch import nn
from torch.optim import Adam

from EnergyConfig import EnergyConfig


class EnergyTrain:

    def __init__(self, model, train_loader, test_loader):
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.model = model
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = Adam(self.model.parameters(), lr=EnergyConfig.learning_rate)

    def train(self):
        """
        Model Train Model
        :return:
        """
        best_acc = 0
        for epoch in range(EnergyConfig.num_epochs):
            self.model.train()
            train_loss = 0
            for inputs, labels in self.train_loader:
                inputs, labels = inputs.to(EnergyConfig.device), labels.to(EnergyConfig.device)
                outputs = self.model(inputs)
                labels = labels.reshape(-1)
                outputs = outputs.reshape(-1)
                loss = self.criterion(outputs, labels)

                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

                train_loss += loss.item() * inputs.size(0)

            # verify
            self.model.eval()
            correct = 0
            total = 0
            with torch.no_grad():
                for inputs, labels in self.test_loader:
                    inputs, labels = inputs.to(EnergyConfig.device), labels.to(EnergyConfig.device)
                    outputs = self.model(inputs)
                    _, predicted = torch.max(outputs.data, dim=0)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()
            train_loss = train_loss / len(self.train_loader.dataset)
            test_acc = correct / total

            print(f'Epoch [{epoch + 1}/{EnergyConfig.num_epochs}] | '
                  f'Train Loss: {train_loss:.4f} | '
                  f'Test Acc: {test_acc:.4f}')

            if test_acc > best_acc:
                best_acc = test_acc
                torch.save(self.model.state_dict(), 'best_model.pth')

        print(f'Best Test Accuracy: {best_acc:.4f}')
